# 选股策略函数，设置各种选股条件
# by yewentao
# v0.1
import time
import _data_collection as dc
import tushare as ts
import _params
import pandas as pd
from os import path

# 使用fama模型，考虑两个因素：PB和market value
# 获取两者乘积最小的股票列表
def fama_with_pb_mv(listNum, tradeDate=''):
    if tradeDate == '':
        today = time.strftime('%Y%m%d')
        trade_date = dc.collect_newest_tradedate(today)
    else:
        trade_date = tradeDate
    ts.set_token(_params.__TL_TOKEN)
    market = ts.Market()
    # 获取当日行情
    stock_hq = market.MktEqud(tradeDate = trade_date)
    if stock_hq.empty:
        print("获取行情数据失败，请检查参数：" + str(tradeDate))
        exit()
    # ticker转成6位字符串
    stock_hq['ticker'] = stock_hq['ticker'].apply(lambda x : str(x).zfill(6))
    # 删除停牌的股票
    stock_hq = stock_hq[stock_hq['isOpen'] == 1]
    # 删除PB为负的
    stock_hq = stock_hq[stock_hq['PB'] > 0]
    # 删除B股
    stock_hq = stock_hq[stock_hq['secID'].apply(lambda x : not x.startswith('900'))]
    # 计算PB*MV
    stock_hq['fama'] = stock_hq['PB'] * stock_hq['marketValue']
    stock_hq = stock_hq.sort_values(by=['fama'])
    return stock_hq.head(listNum)


def kdj_for_buy_signal(kdjN, kdjM1, kdjM2, saveFlag=True):
    # 读取股票列表数据
    stock_list = pd.read_excel(_params.__STOCK_LIST_FULL_FILE_PATH, converters={'ticker':str})
    # 所需要的数据字段
    HQ_FIELDS = ['secID', 'ticker', 'tradeDate', 'closePrice', 'isOpen','highestPrice', 'lowestPrice']
    # 保存金叉信号股票
    golden_list = pd.DataFrame()
    # 循环计算每只股票KDJ
    for ticker in stock_list['ticker']:
        print("开始处理数据：code:" + ticker)
        # 判断是否存在数据文件
        if not path.exists(_params.__STOCK_HQ_FILE_PATH + ticker + ".xlsx"):
            print("数据文件不存在，code:" + ticker)
            continue

        # 读取股票日线数据
        stock_hq = pd.read_excel(_params.__STOCK_HQ_FILE_PATH + ticker + ".xlsx", converters={'ticker':str}, parse_dates=['tradeDate'])
        stock_hq = stock_hq[HQ_FIELDS]
        stock_hq.set_index('tradeDate', inplace=True)
        # 计算N日内的high和low，需要滚动计算
        stock_hq['lown'] = pd.rolling_min(stock_hq['lowestPrice'], kdjN)
        stock_hq['lown'].fillna(value=pd.expanding_min(stock_hq['lowestPrice']), inplace=True)
        stock_hq['highn'] = pd.rolling_max(stock_hq['highestPrice'], kdjN)
        stock_hq['highn'].fillna(value=pd.expanding_max(stock_hq['highestPrice']), inplace=True)
        # 计算N日的RSV（未成熟随机值）数据
        stock_hq['rsv'] = (stock_hq['closePrice'] - stock_hq['lown']) / (stock_hq['highn'] - stock_hq['lown']) * 100
        # 计算K值
        stock_hq['kdj_k'] = pd.ewma(stock_hq['rsv'], kdjM1)
        # 计算D值
        stock_hq['kdj_d'] = pd.ewma(stock_hq['kdj_k'], kdjM2)
        # 计算J值
        stock_hq['kdj_j'] = 3 * stock_hq['kdj_k'] - 2 * stock_hq['kdj_d']

        # 判断KDJ交叉节点
        kd_position = stock_hq['kdj_k'] > stock_hq['kdj_d']
        stock_hq['kdj_gd'] = 0
        # 金叉
        stock_hq.loc[kd_position[(kd_position == True) & (kd_position.shift() == False)].index, 'kdj_gd'] = 1
        # 死叉
        stock_hq.loc[kd_position[(kd_position == False) & (kd_position.shift() == True)].index, 'kdj_gd'] = -1
        # 是否需要保存
        if saveFlag == True:
            stock_hq.to_excel("analysis/kdj/" + ticker + ".xlsx")

        # 分析是否背离，只考虑金叉（最近5个工作日是否有金叉）
        stock_hq_head = stock_hq.tail(5)
        stock_hq_head = stock_hq_head[stock_hq_head['kdj_gd'] == 1]
        if stock_hq_head.empty:
            continue
        # 如果有，考虑背离
        golden_hq = stock_hq[stock_hq['kdj_gd'] == 1]
        first = golden_hq.iloc[len(golden_hq) - 1]
        second = golden_hq.iloc[len(golden_hq) - 2]
        if (first['kdj_j'] > second['kdj_j']) & (first['closePrice'] < second['closePrice']):
            # 出现背离
            golden_list.append(first)

    return golden_list